import os
import re
import time
import subprocess

from threading import Thread

import condor_config as config

S_STATUS_FAILED = -1
S_STATUS_DONE = 0
S_STATUS_RUNNING = 1

POOL_DIR = config.POOL_DIR
STATIC_DIR = config.STATIC_DIR
WAIT_TIME = config.WAIT_TIME


class Condor:
    def __init__(self, jobs, collect, debug=False):
        '''Condor job worker

        Param:
          sched a scheduler implementation

        e.g.
            sched = SomeScheduler(...)
            condor = Condor(sched)

        '''
        if not os.path.exists(POOL_DIR):
            os.mkdir(POOL_DIR)
        self.sched = Scheduler(jobs, collect, debug)
    def run(self):
        while self.sched.is_running():
            mon, job = self.sched.next_job()
            if [mon, job] == [None, None]:
                break

            work_dir = mon.work_dir()
            if not os.path.exists(work_dir):
                os.mkdir(work_dir)

            content = ''.join(open('%s/template_submit'%STATIC_DIR).readlines())
            content = re.sub('XXXXXX', job, content)
            open('%s/submit.txt'%work_dir, 'w+').write(content)
            current_dir = os.getcwd()
            os.chdir(work_dir)
            subprocess.call('condor_submit submit.txt', shell=True)
            os.chdir(current_dir)

            mon.start()



class Scheduler:
    def __init__(self, jobs, collect, debug=False):
        '''Scheduler schedule the job assignment over Condor machines

        Param:
            jobs is a list of parameters in string forms to Condor machines
                 e.g. ['MyJob.py param1 param2', 'MyJob.py param1 param2', ...]
            collect is a customized method to collect finished jobs
                 Collect the results from the monitor, if the job is successful.
                 No need for putting the job back if it is failed.
                 No need for removal of the monitor in the monitor list.

                 True or False is return for the successful/failed result of
                 the collection.

                 def collect(monitor)
        '''
        self.status = S_STATUS_RUNNING
        self.free_machines = range(config.MAX_MACHINE_NUMBER) # free machine ids
        self.next_new_machine = config.MAX_MACHINE_NUMBER
        self.jobs = jobs
        self.collect = collect
        self.monitors = []
        self.debug = debug

    def new_machine(self):
        self.free_machines.append(self.next_new_machine)
        self.next_new_machine += 1

    def next_job(self):
        while self.is_running():
            if self.debug:
                open('log','w+').write('current: %d, jobs: %d, free: %d'%(len(self.monitors),
                                                         len(self.jobs),
                                                         len(self.free_machines)))
                print 'current: %d, jobs: %d, free: %d'%(len(self.monitors),
                                                         len(self.jobs),
                                                         len(self.free_machines))
            for m in self.monitors:
                # check if the monitored job is done
                if not m.is_running():
                    if m.is_failed():
                        self.jobs.append(m.job)
                    else:
                        if not self.collect(m):
                            self.jobs.append(m.job)
                    self.new_machine()
                    m.clean()
                    self.monitors.remove(m)

            # check if there is any free machine
            if len(self.free_machines) > 0:
                # check if there is any work to do
                if len(self.jobs) + len(self.monitors) == 0:
                    self.finish()
                    break # no job any more!
                else:
                    if len(self.jobs) > 0:
                        machine = self.free_machines.pop() # a free machine id is pop out
                        job = self.jobs.pop()
                        mon = Monitor(machine, job)
                        self.monitors.append(mon)
                        return mon, job
                    else:
                        time.sleep(WAIT_TIME)
            else: # wait for free machines or jobs done
                time.sleep(WAIT_TIME) #TODO
        return None, None

    def finish(self):
        self.status = S_STATUS_DONE
        if self.debug:
            subprocess.call('rm log', shell=True)

    def is_running(self):
        return self.status == S_STATUS_RUNNING


class Monitor(Thread):
    def __init__(self, machine, job):
        '''
        Param:
            machine for machine id
            job for job parameters in string form
        '''
        Thread.__init__(self)
        self.machine = machine
        self.job = job
        self.status = S_STATUS_RUNNING

    def work_dir(self):
        return '%s/MACHINE%d'%(POOL_DIR, self.machine)

    def logfile(self):
        return '%s/log.txt'%self.work_dir()

    def outfile(self):
        return '%s/out.txt'%self.work_dir()

    def run(self):
        # wait initialization time of Condor file setup
        time.sleep(WAIT_TIME)

        # real work for detection
        finished = False
        try:
            while not finished:
                fr = open(self.logfile())
                for line in fr.readlines():
                    if 'Job terminated' in line:
                        finished = True
                fr.close()
                time.sleep(WAIT_TIME)
            self.status = S_STATUS_DONE
        except:
            self.status = S_STATUS_FAILED
        
    def clean(self):
        subprocess.call('rm -r %s'%self.work_dir(), shell=True)

    def is_running(self):
        return self.status == S_STATUS_RUNNING
    
    def is_failed(self):
        return self.status == S_STATUS_FAILED





